#!/bin/bash

###==========================================================================
### Usage: bash run_distribute_train_ascend.sh [OPTIONS]...
### Description:
###     Run distributed train for mm model.
###  
### Example:
### 1. Train models with config. Configs in [] are optional.
### bash run_distribute_dynamic_train.sh  configs/ttvitbaseall32guide.yaml mm vibaseflashatten 8 0
###==========================================================================

if [ $# != 5 ]
then
  echo "Usage: bash run_distribute_dynamic_train.sh [CONFIG_PATH] [EXP_TYPE] [EXP_NAME] [RANK_SIZE] [RANK_START]"
  exit 1
fi

get_real_path(){
  if [ "${1:0:1}" == "/" ]; then
    echo "$1"
  else
    echo "$(realpath -m $PWD/$1)"
  fi
}

export DEPLOY_MODE=0
export EXP_TYPE=$2
export EXP_NAME=$3
RANK_SIZE=$4
CONFIG_PATH=$(get_real_path $1)

echo "CONFIG PATH: $CONFIG_PATH"
echo "EXP_TYPE: $EXP_TYPE"
echo "EXP_NAME: $EXP_NAME"
echo "START DEVICE ID: $RANK_START"

export HCCL_EXEC_TIMEOUT=6000
export HCCL_CONNECT_TIMEOUT=1800
ulimit -n 102400


rm -rf dynamic_train
mkdir dynamic_train
cp ../*.py ./dynamic_train
cp -rf ../models ../misc ../misc_utils ../med_lm_finetune ../configs ../data_utils ../examples ../models_old ../models_torch ./dynamic_train
cd ./dynamic_train
mkdir output
LOG_PATH=output
echo "start training"

# 循环启动8个Worker训练进程
for((i=0;i<8;i++));
do
    export MS_WORKER_NUM=8          # 设置集群中Worker进程数量为8
    export MS_SCHED_HOST=127.0.0.1  # 设置Scheduler IP地址为本地环路地址
    export MS_SCHED_PORT=8118       # 设置Scheduler端口
    export MS_ROLE=MS_WORKER        # 设置启动的进程为MS_WORKER角色
    export MS_NODE_ID=$i                      # 设置进程id，可选
    python3 train.py --config $CONFIG_PATH --exp_type $EXP_TYPE --exp_name $EXP_NAME --output_dir $LOG_PATH  2>&1 &    # 启动训练脚本
done

# 启动1个Scheduler进程
export MS_WORKER_NUM=8              # 设置集群中Worker进程数量为8
export MS_SCHED_HOST=127.0.0.1      # 设置Scheduler IP地址为本地环路地址
export MS_SCHED_PORT=8118           # 设置Scheduler端口
export MS_ROLE=MS_SCHED             # 设置启动的进程为MS_SCHED角色
python3 train.py --config $CONFIG_PATH --exp_type $EXP_TYPE --exp_name $EXP_NAME --output_dir $LOG_PATH  2>&1 &

